from model import LLM
from jinja2 import Template
import json
import re

model = LLM()

title_template_str = """
# 角色：专业文本架构师  
## 任务：根据用户主题生成结构化大纲  
**输入参数**：  
- [主题]：{{title}}  
- [篇幅要求]：目标字数范围（默认：2000-3000字）  
- [风格]：技术型/故事型/学术型等（默认：通用）  

**处理规则**：  
1. 生成三级树状结构：  
   - 一级标题（Part）：3-5个核心模块  
   - 二级标题（Chapter）：每个模块下2-4个子章节  
   - 三级标题（Sub_Title）：每个子章节下具体展开点  
2. 确保逻辑链路：  
   `背景引入→核心论点→论证分层→结论升华`  
3. 输出格式：  
```json
{
  "outline": [
    {
      "part": "一级标题名称",
      "chapters": [
        {
          "chapter": "二级标题名称",
          "sub_titles": ["三级标题1", "三级标题2"...]
        }
      ]
    }
  ],
  "total_sub_titles": N // 三级标题总数
}
"""

part_template_str = """
# 角色：领域内容生成专家  
## 任务：根据大纲节点扩写具体内容  
**输入参数**：  
- [当前Part]：{{part}}  
- [当前Chapter]：{{chapter}}  
- [当前Sub_Title]：{{sub_title}}  
- [上下文]：{{final_content}}（JSON格式）  
- [字数要求]：当前段落目标字数（默认：150-200字）  

**处理规则**：  
1. 内容生成三要素：  
   - **知识准确性**：基于可信数据源展开  
   - **逻辑连贯性**：承上启下衔接上下文  
   - **风格一致性**：延续[学术型/故事型]语体  
2. 避免重复策略：  
   - 使用同义词替换  
   - 转换句式结构  
   - 新增数据/案例支撑  
3. 输出格式：  
```json
{
  "content": "生成的段落文本",
  "next_pointer": { // 自动指向下一个待写节点
    "part": "下一模块名",
    "chapter": "下一章节名",
    "sub_title": "下一子标题名"
  },
  "progress": "当前进度百分比"
}
"""

final_content = []

title_template = Template(title_template_str)
prompt = title_template.render({"title": "人工智能的发展与未来"})
rsp_text, tool_calls = model.chat([
    {'role': 'user', 'content': prompt},
])
# response, = model.chat(prompt, temperature=0.5)
# print(rsp_text)
json_pattern = r'```json\n(.*?)\s*```'
content = re.findall(json_pattern, rsp_text, re.DOTALL)
data = json.loads(content[0])

for part in data['outline']:
    part_template = Template(part_template_str)
    part_title = part['part']
    chapters = part['chapters']
    for chapter in chapters:
        chapter_title = chapter['chapter']
        sub_titles = chapter['sub_titles']
        for sub_title in sub_titles:
            print(f"============{part_title},{chapter_title},{sub_title}=============")
            render_data = {
                "part": part_title,
                "chapter": chapter_title,
                "sub_title": sub_title,
                "final_context": final_content
            }
            prompt = part_template.render(render_data)
            rsp_text, tool_calls = model.chat([
                {'role': 'user', 'content': prompt},
            ])

            final_content.append(rsp_text)
